# 环境要求：
# Python 3.x
# TensorFlow 2.x
# Keras 3.x
# numpy
# matplotlib
# tf2onnx
# onnxsim

import os
import tensorflow as tf
import numpy as np
import keras
import matplotlib.pyplot as plt
import onnx
import onnxsim

def load_mnist_data():
    """
    加载MNIST数据集并进行预处理。
    返回预处理后的训练和测试数据。
    """
    (x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data()
    x_train = (x_train.astype("float32") / 255.0)[..., np.newaxis]
    x_test = (x_test.astype("float32") / 255.0)[..., np.newaxis]
    return (x_train, y_train), (x_test, y_test)

def build_and_train_model(x_train, y_train, epochs=5, batch_size=32):
    """
    构建并训练卷积神经网络模型。
    返回训练后的模型和训练历史。
    """
    model = tf.keras.Sequential([
        keras.layers.InputLayer(input_shape=(28, 28, 1)),
        keras.layers.Conv2D(32, (3, 3), activation="relu"),
        keras.layers.Conv2D(64, (3, 3), activation="relu"),
        keras.layers.MaxPooling2D((2, 2)),
        keras.layers.Dropout(0.25),
        keras.layers.Flatten(),
        keras.layers.Dense(10)
    ])

    model.compile(optimizer="adam", 
                  loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True), 
                  metrics=["accuracy"])
    model.summary()
    # 记录训练历史
    history = model.fit(x_train, y_train, epochs=epochs, batch_size=batch_size, validation_split=0.2)
    return model, history

def evaluate_model(model, x_test, y_test):
    """
    评估模型性能。
    """
    loss, accuracy = model.evaluate(x_test, y_test)
    print(f"测试集上的损失: {loss:.4f}, 准确率: {accuracy:.4f}")

def save_as_savedmodel(model, saved_model_dir):
    """
    将模型保存为 SavedModel 格式。
    """
    model.export(saved_model_dir)  # 使用 export 方法保存为 SavedModel 格式
    print(f"模型已保存为 SavedModel 格式到 {saved_model_dir}")

def convert_to_onnx(saved_model_dir, onnx_model_path):
    """
    将 SavedModel 格式的模型转换为 ONNX 格式并保存。
    """
    command = f"python3 -m tf2onnx.convert --saved-model {saved_model_dir} --output {onnx_model_path} --opset 13"
    os.system(command)
    print(f"模型已导出为 ONNX 格式到 {onnx_model_path}")

def simplify_onnx_model(input_model_path, output_model_path):
    """
    使用 onnx-simplifier 简化 ONNX 模型。

    参数:
        input_model_path (str): 输入的 ONNX 模型路径。
        output_model_path (str): 保存简化后 ONNX 模型的路径。
    
    返回:
        bool: 如果模型成功简化，则返回 True；否则返回 False。
    """
    try:
        # 加载原始的 ONNX 模型
        onnx_model = onnx.load(input_model_path)

        # 简化模型
        simplified_model, _ = onnxsim.simplify(onnx_model)

        # 保存简化后的模型
        onnx.save(simplified_model, output_model_path)

        print(f"简化后的 ONNX 模型已保存到: {output_model_path}")
        return True
    except Exception as e:
        print(f"简化 ONNX 模型时出错: {e}")
        return False
    
def convert_to_tflite(model, tflite_model_path, quantize=False):
    """
    将模型转换为 TensorFlow Lite 格式并保存，支持动态范围量化。
    """
    converter = tf.lite.TFLiteConverter.from_keras_model(model)
    if quantize:
        converter.optimizations = [tf.lite.Optimize.DEFAULT]
    tflite_model = converter.convert()
    with open(tflite_model_path, "wb") as f:
        f.write(tflite_model)
    print(f"TFLite 模型已保存到 {tflite_model_path}")

def plot_training_history(history, output_path="training_history.png"):
    """
    Plot the training and validation loss and accuracy over epochs and save as an image file.
    """
    acc = history.history['accuracy']
    val_acc = history.history['val_accuracy']
    loss = history.history['loss']
    val_loss = history.history['val_loss']
    epochs = range(1, len(acc) + 1)

    plt.figure(figsize=(12, 5))

    # Plot accuracy
    plt.subplot(1, 2, 1)
    plt.plot(epochs, acc, 'b', label='Training Accuracy')
    plt.plot(epochs, val_acc, 'r', label='Validation Accuracy')
    plt.title('Training and Validation Accuracy')
    plt.xlabel('Epochs')
    plt.ylabel('Accuracy')
    plt.legend()

    # Plot loss
    plt.subplot(1, 2, 2)
    plt.plot(epochs, loss, 'b', label='Training Loss')
    plt.plot(epochs, val_loss, 'r', label='Validation Loss')
    plt.title('Training and Validation Loss')
    plt.xlabel('Epochs')
    plt.ylabel('Loss')
    plt.legend()

    # Save the image
    plt.tight_layout()
    plt.savefig(output_path)
    plt.show()
    print(f"Training history plot saved to {output_path}")


# 主函数参数设置
saved_model_dir = "saved_model"
onnx_model_path = "model.onnx"
simplify_onnx_model_path = "simplify_model.onnx"
tflite_model_path = "mnist_model.tflite"
epochs = 10
batch_size = 32
quantize = True
output_path = "training_history.png"

def main():
    """
    主函数，整合所有步骤。
    """
    print("加载MNIST数据集...")
    (x_train, y_train), (x_test, y_test) = load_mnist_data()

    print("构建并训练模型...")
    model, history = build_and_train_model(x_train, y_train, epochs=epochs, batch_size=batch_size)

    print("评估模型...")
    evaluate_model(model, x_test, y_test)

    print("绘制训练过程图像...")
    plot_training_history(history, output_path)

    print("保存为 SavedModel 格式...")
    save_as_savedmodel(model, saved_model_dir)

    print("导出为 ONNX 格式...")
    convert_to_onnx(saved_model_dir, onnx_model_path)

    print("简化 ONNX 模型并保存...")
    simplify_onnx_model(onnx_model_path, simplify_onnx_model_path)

    print("转换为 TFLite 模型...")
    convert_to_tflite(model, tflite_model_path, quantize=quantize)

if __name__ == "__main__":
    main()
